XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python poisson2d.py --data_dir=./data/poisson --model=spinn --equation=poisson2d --nc=256 --seed=111 --lr=0.001 --epochs=5000 --mlp=modified_mlp --n_layers=4 --features=64 --r=128 --out_dim=1 --pos_enc=0 --log_iter=500 --plot_iter=5000